package com.bigdata.kafka;

import cn.hutool.core.date.DatePattern;
import cn.hutool.core.date.DateUtil;
import org.apache.kafka.clients.consumer.ConsumerConfig;
import org.apache.kafka.clients.consumer.ConsumerRecord;
import org.apache.kafka.clients.consumer.ConsumerRecords;
import org.apache.kafka.clients.consumer.KafkaConsumer;
import org.apache.kafka.clients.consumer.OffsetAndMetadata;
import org.apache.kafka.clients.consumer.OffsetAndTimestamp;
import org.apache.kafka.common.PartitionInfo;
import org.apache.kafka.common.TopicPartition;

import java.text.ParseException;
import java.text.SimpleDateFormat;
import java.time.Duration;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.Date;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Properties;

/**
 * @author xin.ding
 * @date 2020/11/16 9:52
 */
public class ConsumerTest {

    private static final String bootstrapServer = "10.243.141.133:9092";
    private static final String topic = "test2";

    public static void main(String[] args) {
//        consumerSubscribe();
//        consumerAssign();

        consumerAssignTimestamp("20210407093113");
    }

    /**
     * 直接通过订阅一个指定分区来消费数据
     * (1)如果该groupId消费者分组下 有消费者提交过offset,则从 当前提交的offset位置开始消费数据
     * (2)如果该groupId消费者分组下 没有有消费者提交过offset,则从 当前log添加的最后位置(也就是数据的末尾)开始消费数据
     */
    public static void consumerSubscribe() {
        Properties properties = new Properties();
        properties.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServer);
        properties.put(ConsumerConfig.GROUP_ID_CONFIG, "g1");
//        earliest ：当各分区下存在已提交的 offset 时，从提交的 offset 开始消费；无提交的 offset 时，从头开始消费。
//        latest ：当各分区下存在已提交的 offset 时，从提交的 offset 开始消费；无提交的 offset 时，消费该分区下新产生的数据。
//        none ：topic 各分区都存在已提交的 offset 时，从 offset 后开始消费；只要有一个分区不存在已提交的offset，则抛出异常
        properties.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest");
//        enable.auto.commit 设置为false，如果消费完数据没有提交已消费数据的offset，则会出现重复消费数据的情况
        properties.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, true);
        properties.put(ConsumerConfig.AUTO_COMMIT_INTERVAL_MS_CONFIG, 1000);
        properties.put(ConsumerConfig.MAX_POLL_RECORDS_CONFIG, 1000);
        properties.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, org.apache.kafka.common.serialization.StringDeserializer.class.getName());
        properties.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, org.apache.kafka.common.serialization.StringDeserializer.class.getName());

        final KafkaConsumer<String, String> consumer = new KafkaConsumer<>(properties);
        //订阅topic
        consumer.subscribe(Arrays.asList(topic));
        List<PartitionInfo> parList = consumer.partitionsFor(topic);

        //打印出分区信息
        printPartition(parList);

        //消费数据
        while (true) {
            ConsumerRecords<String, String> records = consumer.poll(Duration.ofMillis(1000));
            for (ConsumerRecord<String, String> record : records) {
                System.out.println(record.toString());
                //手动提交已消费数据的offset
//                    consumer.commitSync();
            }

        }
    }


    /**
     * 通过assign分配的分区,消费者发生故障 Server端不会触发分区重平衡(即使该消费者共享某个已有的groupId)，每个消费者都是独立工作的
     * 为了避免offset提交冲突，需要确保每个消费者都有唯一的groupId
     * 从指定的分区的开头开始消费数据
     */
    public static void consumerAssign() {
        Properties properties = new Properties();
        properties.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServer);
        properties.put(ConsumerConfig.GROUP_ID_CONFIG, "");
        properties.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest");
        properties.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, true);
        properties.put(ConsumerConfig.AUTO_COMMIT_INTERVAL_MS_CONFIG, 1000);
        properties.put(ConsumerConfig.MAX_POLL_RECORDS_CONFIG, 1000);
        properties.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, org.apache.kafka.common.serialization.StringDeserializer.class.getName());
        properties.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, org.apache.kafka.common.serialization.StringDeserializer.class.getName());
        KafkaConsumer<String, String> consumer = new KafkaConsumer<>(properties);
        //获得topic的所有分区
        List<PartitionInfo> partitionList = consumer.partitionsFor(topic);
        //打印出分区信息
        printPartition(partitionList);

        List<TopicPartition> topicPartitions = new ArrayList<>();
        for (PartitionInfo par : partitionList) {
            TopicPartition partition = new TopicPartition(topic, par.partition());
            topicPartitions.add(partition);
        }
        //消费者指定要消费的分区，指定分区之后消费者崩溃之后 不会引发分区reblance
        consumer.assign(topicPartitions);
        //从list中所有分区的开头开始消费数据，这个操作不改变已提交的消费数据的offset
        consumer.seekToBeginning(topicPartitions);

        while (true) {
            ConsumerRecords<String, String> records = consumer.poll(Duration.ofMillis(1000));
            for (ConsumerRecord<String, String> record : records) {
                System.out.println(record.toString());
            }
        }
    }

    /**
     * 通过assign分配的分区,消费者发生故障 Server端不会触发分区重平衡(即使该消费者共享某个已有的groupId)，每个消费者都是独立工作的
     * 为了避免offset提交冲突，需要确保每个消费者都有唯一的groupId
     * 从指定的分区的开头开始消费数据
     */
    public static void consumerAssignTimestamp(String startTime) {
        SimpleDateFormat timeFormat = new SimpleDateFormat("yyyyMMddHHmmss");
        long startTimestamp = 0;
        try {
            Date startDate = timeFormat.parse(startTime);
            startTimestamp = startDate.getTime();
        } catch (ParseException e) {
            e.printStackTrace();
        }

        Properties properties = new Properties();
        properties.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServer);
        properties.put(ConsumerConfig.GROUP_ID_CONFIG, "");
        properties.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest");
        properties.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, true);
        properties.put(ConsumerConfig.AUTO_COMMIT_INTERVAL_MS_CONFIG, 1000);
        properties.put(ConsumerConfig.MAX_POLL_RECORDS_CONFIG, 1000);
        properties.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, org.apache.kafka.common.serialization.StringDeserializer.class.getName());
        properties.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, org.apache.kafka.common.serialization.StringDeserializer.class.getName());
        KafkaConsumer<String, String> consumer = new KafkaConsumer<>(properties);
        //获得topic的所有分区
        List<PartitionInfo> partitionList = consumer.partitionsFor(topic);
        List<TopicPartition> topicPartitions = new ArrayList<>(partitionList.size());
        Map<TopicPartition, Long> timestampsToSearch = new HashMap<>(partitionList.size());

        for (PartitionInfo partitionInfo : partitionList) {
            TopicPartition topicPartition = new TopicPartition(partitionInfo.topic(), partitionInfo.partition());
            topicPartitions.add(topicPartition);
            timestampsToSearch.put(topicPartition, startTimestamp);
        }

        Map<TopicPartition, OffsetAndTimestamp> partitionOffsetAndTimestamp = consumer.offsetsForTimes(timestampsToSearch);
        consumer.assign(new ArrayList<>(timestampsToSearch.keySet()));
        //计算各分区的偏移量
        for (Map.Entry<TopicPartition, OffsetAndTimestamp> entry : partitionOffsetAndTimestamp.entrySet()) {
            // 如果设置的查询偏移量的时间点大于最大的索引记录时间，那么value就为空(取最大索引记录时间对应的偏移量)
            OffsetAndTimestamp offsetTimestamp = entry.getValue();
            long offset;
            if (offsetTimestamp != null) {
                offset = offsetTimestamp.offset();
                final long timestamp = offsetTimestamp.timestamp();
                final String format = DateUtil.format(new Date(timestamp), DatePattern.PURE_DATETIME_PATTERN);
                System.out.println(String.format("start time: %s,date time: %s, offset: %s", startTime, format, offset));

            } else {
                Map<TopicPartition, Long> endOffsets = consumer.endOffsets(topicPartitions);
                offset = endOffsets.get(entry.getKey());
                System.out.println(String.format("查询偏移量的时间点大于最大的索引记录时间,取最大offset: %s", offset));
            }
            consumer.seek(entry.getKey(), offset);
        }

        while (true) {
            ConsumerRecords<String, String> records = consumer.poll(Duration.ofMillis(1000));
            for (ConsumerRecord<String, String> record : records) {
                System.out.println(record.toString());
            }
        }
    }

    public static Long getPartitionLastOffset(KafkaConsumer consumer, Collection<TopicPartition> partitions, TopicPartition topicPartition) {
        try {
            final Map<TopicPartition, Long> endOffsets = consumer.endOffsets(partitions);
            if (endOffsets.containsKey(topicPartition)) {
                return endOffsets.get(topicPartition);
            }

        } catch (Exception e) {
            e.printStackTrace();
        }
        throw new RuntimeException();

    }


    public static void printPartition(List<PartitionInfo> parList) {
        for (PartitionInfo p : parList) {
            System.out.println(p.toString());
        }
    }

    /**
     * 单独处理每个分区中的数据，处理完了之后异步提交offset,注意提交的offset是程序将要读取的下一条消息的offset
     *
     * @param consumer
     */
    public void handlerData(KafkaConsumer<String, String> consumer) {
        boolean running = true;
        try {
            while (running) {
                ConsumerRecords<String, String> records = consumer.poll(Long.MAX_VALUE);
                for (TopicPartition partition : records.partitions()) {
                    List<ConsumerRecord<String, String>> partitionRecords = records.records(partition);
                    for (ConsumerRecord<String, String> record : partitionRecords) {
                        System.out.println(record.offset() + ": " + record.value());
                    }
                    long lastOffset = partitionRecords.get(partitionRecords.size() - 1).offset();
                    //注意提交的offset是程序将要读取的下一条消息的offset
                    consumer.commitSync(Collections.singletonMap(partition, new OffsetAndMetadata(lastOffset + 1)));
                }
            }
        } finally {
            consumer.close();
        }
    }
}
